import oneflow as torch
import oneflow.nn as nn
from .conv import Conv2DFrontEnd
from oasr.frontend.base import BaseFrontEnd



class StackMultiFrontEnd(BaseFrontEnd):
    def __init__(self, n_front_ends, frontend_params):
        super(StackMultiFrontEnd, self).__init__()

        self.ends = nn.ModuleList(
            [
                BuildFrontEnd['front%d_type' % i](**frontend_params['front%d' % i]) for i in range(n_front_ends)
            ]
        )

    def forward(self, inputs, inputs_mask):

        for _, frontend in enumerate(self.ends):
            inputs, inputs_mask = frontend(inputs, inputs_mask)

        return inputs, inputs_mask

    def inference(self, inputs, inputs_mask, cache=None):
        raise NotImplementedError


BuildFrontEnd = {
    'conv': Conv2DFrontEnd,
    'conv2d': Conv2DFrontEnd,
    'deep-conv2d': Conv2DFrontEnd,
    'stack': StackMultiFrontEnd
}